"""
图像生成业务服务
整合提示词提炼、ComfyUI调用和数据库操作
提供完整的图像生成业务流程管理
"""

import asyncio
import logging
import uuid
import json
from typing import Dict, Any, Optional
from datetime import datetime

from ..core.database import get_async_db_session
from ..models.image_generation_task import ImageGenerationTask, TaskStatus
from ..models.visitor import Visitor
from ..models.ai_person import AiPerson
from .prompt_refiner_service import prompt_refiner_service
from .comfyui_service import comfyui_service
from ..core.rate_limiter import get_client_ip, check_rate_limit
from sqlalchemy import select

logger = logging.getLogger(__name__)


class ImageGenerationService:
    """
    图像生成业务服务
    负责协调整个图像生成流程
    """
    
    def __init__(self):
        self.max_concurrent_tasks = 10  # 最大并发任务数
        self.task_timeout = 600  # 任务超时时间（秒）
        
        # 运行中的任务追踪
        self._running_tasks = {}
    
    async def create_image_task(
        self,
        user_uuid: str,
        character_trigger: Optional[str],
        user_prompt: str,
        ai_person_id: Optional[int] = None,
        client_ip: Optional[str] = None,
        user_agent: Optional[str] = None,
        style_name: str = "anime_style_1"
    ) -> Dict[str, Any]:
        """
        创建图像生成任务
        
        Args:
            user_uuid: 用户UUID
            character_trigger: 角色触发词
            user_prompt: 用户原始提示词
            ai_person_id: 关联的AI角色ID（可选）
            client_ip: 客户端IP
            user_agent: 用户代理字符串
            style_name: 工作流样式名称
            
        Returns:
            Dict: 任务创建结果
        """
        try:
            # 验证用户
            async with get_async_db_session() as session:
                result = await session.execute(select(Visitor).filter(Visitor.uuid == user_uuid))
                user = result.scalars().first()
                
                if not user:
                    return {"error": "用户不存在", "code": 404}
                
                if not user.is_active():
                    return {"error": "用户账户已被禁用", "code": 403}
            
            # 检查并发限制
            if len(self._running_tasks) >= self.max_concurrent_tasks:
                return {"error": "服务器繁忙，请稍后再试", "code": 503}
            
            # 创建任务记录
            task_id = uuid.uuid4()
            
            async with get_async_db_session() as session:
                task = ImageGenerationTask(
                    task_id=task_id,
                    user_uuid=user_uuid,
                    ai_person_id=ai_person_id,
                    character_trigger=character_trigger,
                    user_prompt=user_prompt,
                    status=TaskStatus.PENDING.value,
                    client_ip=client_ip,
                    user_agent=user_agent
                )
                
                session.add(task)
                await session.commit()
            
            # 异步开始处理任务
            asyncio.create_task(self._process_task(str(task_id)))
            
            return {
                "task_id": str(task_id),
                "status": TaskStatus.PENDING.value,
                "message": "任务已创建，正在处理中...",
                "estimated_time": 120  # 预估2分钟
            }
            
        except Exception as e:
            logger.error(f"创建图像生成任务失败: {e}", exc_info=True)
            return {"error": "服务器内部错误", "code": 500}
    
    async def _process_task(self, task_id: str):
        """
        异步处理图像生成任务
        
        Args:
            task_id: 任务ID
        """
        try:
            # 添加到运行任务列表
            self._running_tasks[task_id] = datetime.utcnow()
            
            # 更新任务状态为处理中
            await self._update_task_status(task_id, TaskStatus.PROCESSING.value)
            
            # 步骤1: 提炼提示词
            await self._update_task_status(task_id, TaskStatus.PROMPT_REFINING.value)
            refined_prompt = await self._refine_prompt(task_id)
            
            if not refined_prompt:
                await self._update_task_status(task_id, TaskStatus.FAILED.value, "提示词提炼失败")
                return
            
            # 步骤2: 提交到ComfyUI
            await self._update_task_status(task_id, TaskStatus.COMFYUI_QUEUED.value)
            comfyui_prompt_id = await self._submit_to_comfyui(task_id, refined_prompt)
            
            if not comfyui_prompt_id:
                await self._update_task_status(task_id, TaskStatus.FAILED.value, "ComfyUI提交失败")
                return
            
            # 步骤3: 监控ComfyUI任务
            if comfyui_prompt_id == "queued":
                # 任务在排队中，等待后续处理
                logger.info(f"任务 {task_id} 在ComfyUI排队中，等待后续处理")
                return
            else:
                await self._monitor_comfyui_task(task_id, comfyui_prompt_id)
            
        except Exception as e:
            logger.error(f"处理图像生成任务异常: {task_id} - {e}", exc_info=True)
            await self._update_task_status(task_id, TaskStatus.FAILED.value, f"处理异常: {str(e)}")
        finally:
            # 从运行任务列表中移除
            self._running_tasks.pop(task_id, None)
    
    async def _refine_prompt(self, task_id: str) -> Optional[str]:
        """
        提炼任务的提示词
        
        Args:
            task_id: 任务ID
            
        Returns:
            Optional[str]: 提炼后的提示词
        """
        try:
            async with get_async_db_session() as session:
                result = await session.execute(
                    select(ImageGenerationTask).filter(ImageGenerationTask.task_id == task_id)
                )
                task = result.scalars().first()
                
                if not task:
                    logger.error(f"任务不存在: {task_id}")
                    return None
                
                user_prompt = task.user_prompt
            
            # 调用提示词提炼服务
            start_time = datetime.utcnow()
            refined_prompt = await prompt_refiner_service.refine_prompt(user_prompt)
            refining_time = (datetime.utcnow() - start_time).total_seconds()
            
            # 更新数据库
            async with get_async_db_session() as session:
                result = await session.execute(
                    select(ImageGenerationTask).filter(ImageGenerationTask.task_id == task_id)
                )
                task = result.scalars().first()
                
                if task:
                    task.refined_prompt = refined_prompt
                    task.prompt_refining_time = refining_time
                    await session.commit()
            
            # 保留向模型传递的数据日志
                    # 移除详细的提示词日志打印
            
            return refined_prompt
            
        except Exception as e:
            logger.error(f"提示词提炼失败: {task_id} - {e}")
            return None
    
    async def _submit_to_comfyui(self, task_id: str, refined_prompt: str) -> Optional[str]:
        """
        提交任务到ComfyUI
        
        Args:
            task_id: 任务ID
            refined_prompt: 提炼后的提示词
            
        Returns:
            Optional[str]: ComfyUI任务ID
        """
        try:
            async with get_async_db_session() as session:
                result = await session.execute(
                    select(ImageGenerationTask).filter(ImageGenerationTask.task_id == task_id)
                )
                task = result.scalars().first()
                
                if not task:
                    logger.error(f"任务不存在: {task_id}")
                    return None
                
                character_trigger = task.character_trigger
                ai_person_id = task.ai_person_id
                
                # 重新获取角色配置以获取工作流样式
                style_name = "anime_style_1"  # 默认值
                if ai_person_id:
                    try:
                        from .character_service import load_character_profile
                        profile = await load_character_profile(chat_id=ai_person_id)
                        if profile and profile.get('work_flow'):
                            style_name = profile['work_flow']
                    except Exception as e:
        
                        style_name = "anime_style_1"
            
            # 提交到ComfyUI
            comfyui_prompt_id = await comfyui_service.submit_workflow(
                character_trigger=character_trigger,
                refined_prompt=refined_prompt,
                style_name=style_name
            )
            
            if comfyui_prompt_id == "queued":
                # 任务已进入排队，更新状态
                await self._update_task_status(task_id, TaskStatus.COMFYUI_QUEUED.value)
                logger.info(f"任务 {task_id} 已进入ComfyUI排队")
                return "queued"
            elif comfyui_prompt_id:
                # 更新数据库
                async with get_async_db_session() as session:
                    result = await session.execute(
                        select(ImageGenerationTask).filter(ImageGenerationTask.task_id == task_id)
                    )
                    task = result.scalars().first()
                    
                    if task:
                        task.comfyui_prompt_id = comfyui_prompt_id
                        await session.commit()
                
                logger.info(f"任务 {task_id} 已成功提交到ComfyUI，prompt_id: {comfyui_prompt_id}")
            
            return comfyui_prompt_id
            
        except Exception as e:
            logger.error(f"ComfyUI提交失败: {task_id} - {e}")
            return None
    
    async def _monitor_comfyui_task(self, task_id: str, comfyui_prompt_id: str):
        """
        监控ComfyUI任务进度
        
        Args:
            task_id: 任务ID
            comfyui_prompt_id: ComfyUI任务ID
        """
        try:
            max_attempts = 60  # 最大尝试次数
            attempt = 0
            wait_time = 5  # 初始等待时间（秒）
            
            while attempt < max_attempts:
                # 检查任务状态
                status = await comfyui_service.get_task_status(comfyui_prompt_id)
                
                if status.get("status") == "completed":
                    # 移除ComfyUI任务完成日志
                    
                    # 获取生成的图像信息（已包含在status中）
                    image_filename = status.get("image_filename")
                    image_url = status.get("image_url")
                    
                    if image_filename and image_url:
                        # 更新任务为完成状态
                        async with get_async_db_session() as session:
                            result = await session.execute(
                                select(ImageGenerationTask).filter(ImageGenerationTask.task_id == task_id)
                            )
                            task = result.scalars().first()
                            
                            if task:
                                task.image_filename = image_filename
                                task.image_url = image_url
                                task.update_status(TaskStatus.COMPLETED.value)
                                await session.commit()
                        
                        # 移除图像生成完成日志
                        return
                    else:
                        await self._update_task_status(task_id, TaskStatus.FAILED.value, "获取图像结果失败")
                        return
                
                # 处理中，更新进度
                if "progress" in status:
                    progress = status["progress"]
    
                    # 这里可以更新任务进度，但因为进度变化频繁，可能不需要每次都写入数据库
                
                # 等待一段时间再次检查
                await asyncio.sleep(wait_time)
                attempt += 1
                
                # 指数退避策略：逐渐增加等待时间
                if attempt % 5 == 0 and wait_time < 30:
                    wait_time += 5
            
            # 超过最大尝试次数
            logger.error(f"ComfyUI任务监控超时: {task_id}")
            await self._update_task_status(task_id, TaskStatus.FAILED.value, "任务处理超时")
            
        except Exception as e:
            logger.error(f"监控ComfyUI任务异常: {task_id} - {e}")
            await self._update_task_status(task_id, TaskStatus.FAILED.value, f"监控异常: {str(e)}")
    
    async def _update_task_status(self, task_id: str, status: str, error_message: str = None):
        """
        更新任务状态
        
        Args:
            task_id: 任务ID
            status: 新状态
            error_message: 错误信息（可选）
        """
        try:
            async with get_async_db_session() as session:
                result = await session.execute(
                    select(ImageGenerationTask).filter(ImageGenerationTask.task_id == task_id)
                )
                task = result.scalars().first()
                
                if task:
                    task.update_status(status, error_message)
                    await session.commit()
        
                
        except Exception as e:
            logger.error(f"更新任务状态失败: {task_id} - {e}")
    
    async def get_task_status(self, task_id: str, user_uuid: str) -> Dict[str, Any]:
        """
        获取任务状态
        
        Args:
            task_id: 任务ID
            user_uuid: 用户UUID
            
        Returns:
            Dict: 任务状态信息
        """
        try:
            async with get_async_db_session() as session:
                result = await session.execute(
                    select(ImageGenerationTask).filter(
                        ImageGenerationTask.task_id == task_id,
                        ImageGenerationTask.user_uuid == user_uuid
                    )
                )
                task = result.scalars().first()
                
                if not task:
                    return {"error": "任务不存在", "code": 404}
                
                return task.to_dict()
                
        except Exception as e:
            logger.error(f"获取任务状态失败: {task_id} - {e}")
            return {"error": "服务器内部错误", "code": 500}
    
    async def cancel_task(self, task_id: str, user_uuid: str) -> Dict[str, Any]:
        """
        取消图像生成任务
        
        Args:
            task_id: 任务ID
            user_uuid: 用户UUID
            
        Returns:
            Dict: 取消结果
        """
        try:
            async with get_async_db_session() as session:
                result = await session.execute(
                    select(ImageGenerationTask).filter(
                        ImageGenerationTask.task_id == task_id,
                        ImageGenerationTask.user_uuid == user_uuid
                    )
                )
                task = result.scalars().first()
                
                if not task:
                    return {"error": "任务不存在", "code": 404}
                
                # 检查状态是否可以取消
                if task.status in [TaskStatus.COMPLETED.value, TaskStatus.FAILED.value, TaskStatus.CANCELLED.value]:
                    return {
                        "message": f"任务已处于{task.status}状态，无需取消",
                        "status": task.status
                    }
                
                # 如果有ComfyUI任务ID，尝试取消ComfyUI任务
                if task.comfyui_prompt_id:
                    try:
                        await comfyui_service.cancel_workflow(task.comfyui_prompt_id)
                    except Exception as e:
                        pass
                
                # 更新任务状态
                task.update_status(TaskStatus.CANCELLED.value)
                await session.commit()
                
                # 从运行任务列表中移除
                self._running_tasks.pop(str(task.task_id), None)
                
                return {
                    "message": "任务已取消",
                    "status": TaskStatus.CANCELLED.value
                }
                
        except Exception as e:
            logger.error(f"取消任务失败: {task_id} - {e}")
            return {"error": "服务器内部错误", "code": 500}
    
    async def get_user_tasks(
        self,
        user_uuid: str,
        limit: int = 20,
        offset: int = 0
    ) -> Dict[str, Any]:
        """
        获取用户的图像生成任务列表
        
        Args:
            user_uuid: 用户UUID
            limit: 返回的最大任务数
            offset: 分页偏移量
            
        Returns:
            Dict: 用户任务列表和分页信息
        """
        try:
            async with get_async_db_session() as session:
                # 查询任务总数
                result_count = await session.execute(
                    select(ImageGenerationTask).filter(
                        ImageGenerationTask.user_uuid == user_uuid
                    )
                )
                total = len(result_count.scalars().all())
                
                # 查询任务列表
                result = await session.execute(
                    select(ImageGenerationTask).filter(
                        ImageGenerationTask.user_uuid == user_uuid
                    ).order_by(ImageGenerationTask.created_at.desc()).offset(offset).limit(limit)
                )
                tasks = result.scalars().all()
                
                task_list = [task.to_dict() for task in tasks]
                
                return {
                    "tasks": task_list,
                    "total": total,
                    "limit": limit,
                    "offset": offset
                }
                
        except Exception as e:
            logger.error(f"获取用户任务列表失败: {user_uuid} - {e}")
            return {"error": "服务器内部错误", "code": 500}
    
    async def get_service_status(self) -> Dict[str, Any]:
        """
        获取服务状态
        
        Returns:
            Dict: 服务状态信息
        """
        return {
            "active": True,
            "running_tasks": len(self._running_tasks),
            "max_concurrent_tasks": self.max_concurrent_tasks,
            "available_slots": max(0, self.max_concurrent_tasks - len(self._running_tasks))
        }


# 创建全局服务实例
image_generation_service = ImageGenerationService() 